import torch

torch.manual_seed(20)  # 为CPU设置随机种子
torch.cuda.manual_seed(20)  # 为当前GPU设置随机种子
torch.cuda.manual_seed_all(20)  # 为所有GPU设置随机种子
# 生成一个每个元素均为[100-150]均匀分布的4行3列随机张量
x = torch.randint(100, 150, (2, 3, 4, 5))
x = torch.as_tensor(x, dtype=torch.float32)
print("x:", x)
y = torch.zeros_like(x)  # 输出
# 参数准备
num_features = channel = x.shape[1]
momentum = 0.1
eps = 1e-5

# 初始化参数
running_mean = torch.zeros(channel)
running_var = torch.ones(channel)
gamma = torch.ones(channel)
beta = torch.zeros(channel)

# 一.train 模式
# 1. 计算 saved_mean,saved_var
saved_mean = x.mean(dim=(0, 2, 3))  # (3,)计算 3个通道的均值
# 或：
# saved_mean = x.mean(dim=0, keepdim=True).
#                mean(dim=2, keepdim=True).
#                mean(dim=3, keepdim=True)
# keepdim=True保持维度不变

saved_var = x.var(dim=(0, 2, 3), unbiased=False)
# 计算3个通道方差
# unbiased参数设置为False,否则计算出方差为无偏估计,见代码解释3
print("saved_mean:", saved_mean)
print("saved_var:", saved_var)
# 2. BN 前向
for c in range(channel):
    y[:, c, :, :] = ((x[:, c, :, :] - saved_mean[c]) /
                     torch.pow(saved_var[c] + eps, 0.5)) * gamma[c] + beta[c]
print(y)

# 3. 更新全局数据均值running_mean,全局数据方差running_var
# 通过有偏方差和无偏方差的关系,又转换成了无偏方差,unbiased设为True
saved_variance = x.var(dim=(0, 2, 3), unbiased=True)

running_mean = (1 – momentum) * running_mean + momentum * saved_mean
running_var = (1 - momentum) * running_var + momentum * saved_variance
print("running_mean:", running_mean)
print("running_var:", running_var)

# 二.eval 模式
#  BN 前向,测试过程中利用存储的均值和方差对各个通道进行标准化处理
for c in range(channel):
    y[:, c, :, :] = ((x[:, c, :, :] - running_mean[c]) /
                     torch.pow(running_var[c] + eps, 0.5)) * gamma[c] + beta[c]
